open Common

(* -------------------------------------------------------------------------- *)
(*  Imperative heaps                                                          *)
(* -------------------------------------------------------------------------- *)

(** Mostly follows CLRS, Chapter 6 (Heaps) and 17.4 (Dynamic tables) *)
module type GEN = 
sig

  module O : Ordered.T 

  exception Empty

  (* Nodes *)
  type 'a node 
  val value : 'a node -> 'a
  val priority : 'a node -> O.t
  val node_to_string : ('a -> string) -> 'a node -> string

  (* DEBUG *)
  val set_prio : 'a node -> O.t -> unit


  (* Heaps *)
  type 'a t
  val create : unit -> 'a t
  val build : 'a node array -> 'a t
  val insert : 'a t -> O.t * 'a -> 'a node
  val heapsort : 'a node array -> unit
  val size : 'a t -> int
  val to_string : ('a -> string) -> 'a t -> string
  val empty : 'a t -> bool

  val grow_factor : int ref
  val shrink_factor : int ref

end

module type IMP_MAX = 
sig
  include GEN
  val maximum : 'a t -> O.t * 'a
  val extract_max : 'a t -> O.t * 'a (* raises Empty *)
  val increase_key : 'a t -> 'a node -> O.t -> unit
end

module type IMP_MIN = 
sig
  include GEN
  val minimum : 'a t -> O.t * 'a
  val extract_min : 'a t -> O.t * 'a (* raises Empty *)
  val decrease_key : 'a t -> 'a node -> O.t -> unit
end

(* -------------------------------------------------------------------------- *)
(*  Default implementation                                                    *)
(* -------------------------------------------------------------------------- *)

type kind = Min_heap | Max_heap

module type KIND = 
sig 
  val kind : kind 
end

module MaxKind : KIND = 
struct
  let kind = Max_heap
end

module MinKind : KIND =
struct 
  let kind = Min_heap
end

module S(O : Ordered.T)(K : KIND) =
struct 

  module A = Array'

  (* --------------------------------  Types  --------------------------------- *)

  type 'a node = {x : 'a;
                  prio : O.t ref;
                  index : int ref}

  type 'a t = {arr : 'a node array ref;
               size : int ref;
               default : 'a option ref}

  exception Empty

  (* --------------------------------  Nodes  --------------------------------- *)

  let mk_node prio x = {x = x;
                        prio = ref prio;
                        index = ref (-1)}
  let value n = n.x
  let priority n = !(n.prio)
  let node_to_string f {x=x; prio=prio; index=index} = "{" ^ f x ^ " : " ^ O.to_string !prio ^ " > " ^ string_of_int !index ^ "}"
  (* DEBUG *)
  let set_prio n y = n.prio := y
  
  (* --------------------------  Dynamic Expansion  --------------------------- *)
  
  let grow_factor = ref 2

  let shrink_factor = ref 4

  let should_grow {arr=arr;size=size} = !size = A.length !arr

  let grow {arr=arr;size=size;default=default} = 
    if is_none !default then failwith "grow: no default" else
    let n = A.length !arr in
    let new_size = if n = 0 then 1 else !grow_factor * n in
    let arr' = A.create new_size (mk_node O.default (the !default)) in
      begin 
        print_endline ("growing heap size to: " ^ string_of_int new_size);
        A.blit !arr 0 arr' 0 !size;
        arr := arr'
      end

  let shrink_ratio = 1.0 /. float_of_int !shrink_factor 
  let should_shrink {arr=arr;size=size} =
      float_of_int !size /. float_of_int (A.length !arr) < shrink_ratio

  let shrink {arr=arr;size=size} = 
    let new_size = A.length !arr / !grow_factor in
    let arr' = A.create new_size !arr.(0) in
      begin 
        A.blit !arr 0 arr' 0 !size;
        arr := arr'
      end

  (* -------------------------------  Heap ops  ------------------------------- *)

  let parent i = if i = 0 then failwith "parent: root" else ((i + 1) / 2) - 1
  let left i = (2 * (i + 1)) - 1 
  let right i = 2 * (i + 1)

  let arr heap = !(heap.arr)
  let node heap n = (arr heap).(n)
  let size heap = !(heap.size)
  let empty heap = size heap = 0
  let prio heap n = !((node heap n).prio)
  let index heap n = !((node heap n).index)

  let hswap heap i j = 
    let ai = node heap i in
    let aj = node heap j in
    begin 
      ai.index := j;
      aj.index := i;
      (arr heap).(j) <- ai;
      (arr heap).(i) <- aj
    end

  let hinsert ({arr=arr;size=size;default=default} as heap) ({prio=prio;index=index;x=x} as node) = 
    begin 
      if !default = None then default := Some x;
      if should_grow heap then grow heap;
      !arr.(!size) <- node;
      index := !size;
      incr size;
    end

  let lt = match K.kind with Max_heap -> O.lt | Min_heap -> O.gt 

  let max_heapify i =
    let rec heapify ({size=size} as heap) i = 
      let ip = prio heap i in
      let largest = ref i in
      let l,r = left i,right i in
      begin 
        if l < !size && lt ip (prio heap l) then largest := l;
        if r < !size && lt (prio heap !largest) (prio heap r) then largest := r;
        if !largest <> i then (hswap heap i !largest; heapify heap !largest)
      end in
      heapify i

  let create() = {arr = ref [||];
                  size = ref 0;
                  default = ref None}
                    
  let size {size=size} = !size

  let build xarr = 
    let n = A.length xarr in
    let ret = {arr = ref xarr; size = ref n; default = ref None} in
    if n = 0 then ret else
    begin 
      A.iteri (fun n x -> x.index := n) xarr;
      let i = ref (n / 2) in
      ret.default := Some (xarr.(0).x);
      while !i >= 0 do
        max_heapify ret !i;
        decr i
      done;
      ret
    end

  let maximum {arr=arr;size=size} = 
    if !size = 0 then raise Empty else 
    let node = !arr.(0) in
      priority node, value node

  let extract_max ({arr=arr;size=size} as heap) = 
    if !size <= 0 then raise Empty else
    let ret = !arr.(0) in
    begin
      hswap heap 0 (!size - 1);
      decr size;
      max_heapify heap 0;
      if should_shrink heap then shrink heap;
      !(ret.prio),ret.x
    end

  let increase_key ({arr=arr} as heap) {prio=prio;index=index} k =
    if lt k !prio then failwith "increase_key: new key smaller than current key" else
    let pprio() = !(!arr.(parent (!index)).prio) in
    begin 
      prio := k;
      while !index > 0 && lt (pprio()) !prio do
        hswap heap (!index) (parent !index);
      done 
    end

  let insert heap (prio,x) =
    let node = mk_node prio x in
    begin 
      hinsert heap node;
      increase_key heap node prio;
      node
    end

  let to_string f {arr=arr;size=size} = 
    let l,_ = chop_list !size (A.to_list !arr) in
      curly (string_of_int !size) ^ (string_of_list (node_to_string f) l)

  let heapsort xarr = 
    let ({arr=arr;size=size} as heap) = build xarr in
    let i = ref (!size - 1) in
    while !i >= 1 do
      hswap heap 0 !i;
      decr i;
      decr size;
      max_heapify heap 0;
    done

  let decrease_key = increase_key
  let minimum = maximum
  let extract_min = extract_max

end 

module MaxHeap(O' : Ordered.BASE) : IMP_MAX with module O = Ordered.S(O') =
struct 
  module O = Ordered.S(O')
  module H = S(O)(MaxKind)
  include H
end

module MinHeap(O' : Ordered.BASE) : IMP_MIN with module O = Ordered.S(O') =
struct 
  module O = Ordered.S(O')
  module H = S(O)(MinKind)
  include H
end

module MaxInt = MaxHeap(Ordered.Int)
module MinInt = MinHeap(Ordered.Int)

(* -------------------------------------------------------------------------- *)
(*  Persistent heaps                                                          *)
(* -------------------------------------------------------------------------- *)

module type PER_MAX =
sig
end

module type PER_MIN =
sig
end
